package com.cl.ks.flow.aspect;

import com.alibaba.fastjson.JSON;

import com.cl.ks.entity.*;
import com.cl.ks.flow.base.BaseFlowNodeHandler;
import com.cl.ks.flow.util.StaticCollection;
import com.cl.ks.repository.KsFlowGraphRepository;
import com.cl.ks.service.KsFlowGraphLogService;
import com.cl.ks.service.KsFlowGraphRecordService;
import com.cl.ks.service.KsFlowNodeResultService;
import com.cl.ks.utils.ExceptionUtil;
import com.cl.spark.dto.SparkResult;
import com.querydsl.jpa.impl.JPAQueryFactory;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.SaveMode;
import org.apache.spark.sql.SparkSession;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.BeanUtils;
import org.springframework.stereotype.Component;

import java.io.File;
import java.io.IOException;
import java.lang.reflect.UndeclaredThrowableException;
import java.nio.file.Paths;
import java.time.LocalDateTime;
import java.util.List;
import java.util.Map;

import static com.cl.ks.flow.base.BaseFlowNodeHandler.newKsFlowNodeResult;
import static org.apache.commons.io.FileUtils.deleteDirectory;


@Aspect
@Component
@Slf4j
public class KsFlowAspect {

    private final KsFlowGraphRecordService ksFlowGraphRecordService;

    private final KsFlowGraphRepository ksFlowGraphRepository;

    private final KsFlowGraphLogService ksFlowGraphLogService;

    private final KsFlowNodeResultService ksFlowNodeResultService;

    private final SparkSession sparkSession;


    public KsFlowAspect(KsFlowGraphRecordService ksFlowGraphRecordService, KsFlowGraphRepository ksFlowGraphRepository, KsFlowGraphLogService ksFlowGraphLogService, KsFlowNodeResultService ksFlowNodeResultService, SparkSession sparkSession) {
        this.ksFlowGraphRecordService = ksFlowGraphRecordService;
        this.ksFlowGraphRepository = ksFlowGraphRepository;
        this.ksFlowGraphLogService = ksFlowGraphLogService;
        this.ksFlowNodeResultService = ksFlowNodeResultService;
        this.sparkSession = sparkSession;
    }

    @Pointcut("execution(public * com.cl.ks.flow.base.BaseFlowNodeHandler.handle(..))")
    public void process() {
    }

    @Pointcut("execution(public * com.cl.ks.flow.KsFlowGraphHandler.handleKsFlowGraph(..))")
    public void handleKsFlowGraph() {
    }

    @Around("process()")
    public SparkResult process(ProceedingJoinPoint joinPoint) {
        SparkResult res;
        BaseFlowNodeHandler.ProcessParam processParam = (BaseFlowNodeHandler.ProcessParam) joinPoint.getArgs()[0];
        KsFlowGraphLog ksFlowGraphLog = processParam.getKsFlowGraphLog();
        SparkNode sparkNode = processParam.getSparkNode();
        KsFlowGraph ksFlowGraph = processParam.getKsFlowGraph();
        KsFlowGraphRecord ksFlowGraphRecord = processParam.getKsFlowGraphRecord();
        ksFlowGraphLog.setSparkNodeCode(sparkNode.getCode());
        ksFlowGraphLog.setKsFlowGraphId(ksFlowGraph.getId());
        ksFlowGraphLog.setKsFlowGraphRecordId(ksFlowGraphRecord.getId());
        SparkNode sparkNodeInfo = new SparkNode();
        BeanUtils.copyProperties(sparkNode, sparkNodeInfo);
        sparkNodeInfo.setInputSparkResultList(null);
        sparkNodeInfo.setNextSparkNodes(null);
        ksFlowGraphLog.setSparkNodeJson(JSON.toJSONString(sparkNodeInfo));
        ksFlowGraphLog.setKsFlowTaskId(ksFlowGraph.getKsFlowTaskId());
        ksFlowGraphLog.setCreatedBy(ksFlowGraph.getCreatedBy());
        ksFlowGraphLog.setCreatedByShowName(ksFlowGraph.getCreatedByShowName());

        ksFlowGraphLogService.save(ksFlowGraphLog);

        String key = "";
        String nodeName = sparkNode.getNodeName();
        if ("模型".equals(nodeName)) {
            key = (String) sparkNode.getNodeExpression().get("ruleName");
        } else if ("数据源".equals(nodeName)) {
            key = (String) sparkNode.getNodeExpression().get("targetTableName");
        } else if (sparkNode.getNodeExpression() != null) {
            key = (String) sparkNode.getNodeExpression().get("name");
        }
        if (StringUtils.isBlank(key)) {
            key = nodeName;
        }
        Map<String, SparkResult> inputSparkResultList = sparkNode.getInputSparkResultList();

        String sparkRddDir = "./spark-rdd/" + ksFlowGraph.getId() + "/";
        if (StringUtils.isNotBlank(ksFlowGraph.getFailedNodeCode()) && ksFlowGraph.getFailedNodeCode().equals(sparkNode.getCode())) {
            inputSparkResultList.forEach((code, sparkResult) -> {
                try {
                    Dataset<Row> checkpointsDataset = sparkSession.read().parquet(sparkRddDir + sanitizePathSegment(code));
                    sparkResult.setDataset(checkpointsDataset);
                    sparkResult.setCount(checkpointsDataset.count());
                } catch (Exception e) {
                    e.printStackTrace();
                }

            });
        }

        long inputCount = 0L;
        if (inputSparkResultList != null) {
            for (SparkResult i : inputSparkResultList.values()) {
                if (i.getCount() != null) {
                    inputCount += i.getCount();
                }
            }
        }

        ksFlowGraphLog.setInputSize(inputCount);
        ksFlowGraphLog.setSparkNodeName(key);

        ksFlowGraphLog.setStatus(1);
        ksFlowGraphLogService.save(ksFlowGraphLog);
//        KsFlowGraphHandler.setProgress(ksFlowGraphLog);
        // proceed
        try {

            try {
                res = (SparkResult) joinPoint.proceed();
            }catch (UndeclaredThrowableException e) {
                throw e.getCause();
            }

            // 存储日志
            ksFlowGraphLog.setOutputSize(res.getCount());
            ksFlowGraphLog.setStatus(10);
            res.setPositionLeft(sparkNode.getPositionLeft());

            // 流程图存储关联数据
            BaseFlowNodeHandler.ProcessParam afterProcessParam = (BaseFlowNodeHandler.ProcessParam) joinPoint.getArgs()[0];
            if (!ksFlowGraph.getSparkNodeJson().contains("inputSparkResultList")) {
                ksFlowGraphRepository.save(afterProcessParam.getKsFlowGraph());
            }
            // 存储节点输出数据
            List<KsFlowNodeResult> ksFlowNodeResultList = processParam.getKsFlowNodeResultList();
            int size = 100;
            if (ksFlowNodeResultList.isEmpty() && res.getDataset() != null) {
                Integer ksFlowGraphRecordId = processParam.getKsFlowGraphRecord().getId();
                Integer ksFlowGraphLogId = processParam.getKsFlowGraphLog().getId();
                Integer ksFlowGraphId = processParam.getKsFlowGraph().getId();
                String thisNodeName = processParam.getSparkNode().getNodeName();
                String nodeCode = processParam.getSparkNode().getCode();
                Integer ksFlowTaskId = processParam.getKsFlowGraph().getKsFlowTaskId();
                String createdBy = processParam.getKsFlowGraph().getCreatedBy();
                String createdByShowName = processParam.getKsFlowGraph().getCreatedByShowName();
                String power = processParam.getKsFlowGraph().getPower();


                if (res.getCount() != null && res.getCount() < 10000) {
                    StaticCollection.newKsFlowNodeResultList(ksFlowGraphId);
                    res.getDataset().limit(size).foreach(row -> {
                        StaticCollection.getKsFlowNodeResultList(ksFlowGraphId).add(newKsFlowNodeResult(
                                ksFlowGraphRecordId,
                                ksFlowGraphLogId,
                                ksFlowGraphId,
                                thisNodeName,
                                nodeCode,
                                row.json(),
                                ksFlowTaskId,
                                createdBy,
                                createdByShowName,
                                power
                        ));
                    });
                    ksFlowNodeResultService.saveAll(StaticCollection.getKsFlowNodeResultList(ksFlowGraphId));
                }
                StaticCollection.removeKsFlowNodeResult(ksFlowGraphId);
            }
            if (sparkNode.getNodeName().equals("结束")) {
                deleteDirectory(new File(sparkRddDir));
                if (StringUtils.isNotBlank(ksFlowGraph.getFailedNodeCode())) {
                    ksFlowGraph.setFailedNodeCode(null);
                    if (!ksFlowGraph.getSparkNodeJson().contains("inputSparkResultList")) {
                        ksFlowGraphRepository.save(ksFlowGraph);
                    }
                }

            }
        } catch (Throwable e) {
            ksFlowGraphRecord.setStatus(-1);
            ksFlowGraphRecordService.save(ksFlowGraphRecord);
            // 输出到控制台或日志系统
            e.printStackTrace();

            // 设置到ksFlowGraphLog中
            ksFlowGraphLog.setMsg(key + ":" + ExceptionUtil.getExceptionChainAsString(e));
            ksFlowGraphLog.setStatus(-1);

            if (inputSparkResultList != null) {
                if (inputSparkResultList.values().stream().anyMatch(item -> item.getCount() > 0)) {
                    inputSparkResultList.forEach((code, sparkResult) -> {
                        if (sparkResult.getCount() > 0) {
                            String thisSparkRddDirPath = sparkRddDir + sanitizePathSegment(code);
                            Dataset<Row> dataset = sparkResult.getDataset();
                            File thisSparkRddDir = new File(thisSparkRddDirPath);
                            if (!thisSparkRddDir.exists()) {
                                thisSparkRddDir.mkdirs();
                                dataset.write().mode(SaveMode.Overwrite).parquet(thisSparkRddDirPath);
                            }
                        }
                    });
                    ksFlowGraph.setFailedNodeCode(sparkNode.getCode());
                    if (!ksFlowGraph.getSparkNodeJson().contains("inputSparkResultList")) {

                        ksFlowGraphRepository.save(ksFlowGraph);
                    }
                }


            }
            throw new RuntimeException(e.getMessage());
        } finally {
            ksFlowGraphLog.setEndTime(LocalDateTime.now());
            try {
                ksFlowGraphLogService.save(ksFlowGraphLog);
            } catch (Exception e) {
                e.printStackTrace();
                if (ksFlowGraphLog.getMsg() != null && ksFlowGraphLog.getMsg().length() > 255) {
                    ksFlowGraphLog.setMsg(ksFlowGraphLog.getMsg().substring(0, 255));
                    ksFlowGraphLogService.save(ksFlowGraphLog);
                }
            }
        }
        return res;
    }

    private static String sanitizePathSegment(String pathSegment) {
        // 移除非法字符
        return pathSegment.replaceAll("[^A-Za-z0-9_.-]", "");
    }
//
//    @Around("handleKsFlowGraph()") //在切入点的方法run之前要干的
//    public Object handleKsFlowGraph(ProceedingJoinPoint joinPoint) {
//        Object result = null;
//        try {
//            KsFlowGraph ksFlowGraph = (KsFlowGraph) joinPoint.getArgs()[0];
//            KsFlowGraphRecord ksFlowGraphRecord = ksFlowGraphRecordService.saveKsFlowGraphRecord(ksFlowGraph);
//
//            result = joinPoint.proceed();
//        } catch (Throwable e) {
//            throw new RuntimeException(e);
//        }
//        return result;
//    }


}
